# -*- coding: utf-8 -*-
"""
   File Name:  plot.py
   Author :    liccoo
   Time:       2022/8/24 12:46
"""
import os
from random import randint

import numpy as np
from matplotlib import pyplot as plt


def create_x_data(start_point, end_point, num_points):   # 创建横坐标
    x_data = np.linspace(start_point, end_point, num_points)
    return x_data


def plot_predict_truth(_x_data, _predict_value, _truth_value, _path, _index):
    """
    Plot the predicted value and true one in a graph

    :param _x_data: ndarray
    :param _predict_value: ndarray
    :param _truth_value: ndarray
    :param _path: ndarray
    :param _index: ndarray
    :return: None
    """
    # 设置图像的大小，宽度为 20 * 10 ，并设置分辨率为 300 ppi
    plt.figure(figsize=(10, 10), dpi=72)

    # 绘制预测值与实际值
    plt.plot(_x_data, _predict_value, 'r', marker='*', label="Predicted")  # 预测值
    plt.plot(_x_data, _truth_value, 'g', marker='v', label="True")  # 真实值
    plt.ylim([-1.1, 1.1])  # 设置 Y坐标轴 范围
    plt.xlabel(' ')
    plt.ylabel(' ')
    plt.title('The comparison of predicted value and true one')
    plt.legend(loc='upper right')

    # 保存并展示图片
    plt.savefig(os.path.join(_path, '%.3d_predicted_true.png' % _index))
    plt.show()


def plot_forward(args, model, dataset, _path, start_point, end_point, num_points):
    # 打印指定的图片数  ->  plot_num
    for index in range(args.plot_num):
        # 判断是否随机抽样
        if args.random_sample:
            dataset_len = len(dataset)
            random_index = randint(0, dataset_len - 1)  # randint 包含两端
        else:
            random_index = index  # 如果不随机抽样，则会按照测试集的顺序采样

        # 加载数据集里的数据  ---  不同的数据集需进行不一样的读取设置
        x = dataset[random_index][0]
        y = dataset[random_index][1]

        # 真实的 y 值
        truth_y = y.detach().numpy()

        # 计算网络预测结果
        predict = model(x)
        predict_y = predict.detach().numpy()  # 网络预测值转换需要加 .detach()

        print('true x: %.4f  predicted x: %.4f' % (truth_y, predict_y))

        # 创建横坐标的值 x_data
        x_data = create_x_data(start_point, end_point, num_points)

        # 开始绘制图形
        plot_predict_truth(
            _x_data=x_data,
            _predict_value=predict_y,
            _truth_value=truth_y,
            _path=_path,
            _index=index
        )


def plot_inverse(args, model, dataset, _path, start_point, end_point, num_points):
    # 打印指定的图片数  ->  plot_num
    for index in range(args.plot_num):
        # 判断是否随机抽样
        if args.random_sample:
            dataset_len = len(dataset)
            random_index = randint(0, dataset_len - 1)  # randint 包含两端
        else:
            random_index = index  # 如果不随机抽样，则会按照测试集的顺序采样

        # 加载数据集里的数据  ---  不同的数据集需进行不一样的读取设置
        y = dataset[random_index][0]
        x = dataset[random_index][1]

        # 真实的 x 值
        truth_x = x.detach().numpy()

        # 计算网络预测结果
        predict = model(y)
        predict_x = predict.detach().numpy()  # 网络预测值转换需要加 .detach()

        print('true x: %.4f  predicted x: %.4f' % (truth_x, predict_x))

        # 创建横坐标的值 x_data
        x_data = create_x_data(start_point, end_point, num_points)

        # 开始绘制图形
        plot_predict_truth(
            _x_data=x_data,
            _predict_value=predict_x,
            _truth_value=truth_x,
            _path=_path,
            _index=index
        )


def plot_inverse_scatter(args, model, dataset, _path):
    # 设置图像的大小，宽度为 10 * 10 ，并设置分辨率为 72 ppi
    plt.figure(figsize=(10, 10), dpi=72)

    # 打印指定的图片数  ->  plot_num
    for index in range(args.plot_num):
        # 判断是否随机抽样
        if args.random_sample:
            dataset_len = len(dataset)
            random_index = randint(0, dataset_len - 1)  # randint 包含两端
        else:
            random_index = index  # 如果不随机抽样，则会按照测试集的顺序采样

        # 加载数据集里的数据  ---  不同的数据集需进行不一样的读取设置
        y = dataset[random_index][0]
        x = dataset[random_index][1]

        # 真实的 x 值
        truth_x = x.detach().numpy()

        # 计算网络预测结果
        predict = model(y)
        predict_x = predict.detach().numpy()  # 网络预测值转换需要加 .detach()

        print('true x: %.4f  predicted x: %.4f' % (truth_x, predict_x))
        true_y = y.detach().numpy()

        # 绘制实部
        plt.scatter(predict_x, true_y, c='r', s=10, label='Predict')
        plt.scatter(truth_x, true_y, c='b', s=10, label='Truth')

        predict_y = np.sin(5 * predict_x) * np.cos(10 * predict_x)
        print('true y: %.6f  predicted y: %.6f' % (true_y, predict_y))

    x_ = np.linspace(-1, 1, 100)
    y_ = np.sin(5 * x_) * np.cos(10 * x_)  # y = math.sin(5 * x) * math.cos(10 * x)
    plt.plot(x_, y_, 'y', marker='*', label='Function')  # 预测值
    plt.xlim([-1.1, 1.1])  # 设置 X坐标轴 范围
    plt.ylim([-1.1, 1.1])  # 设置 Y坐标轴 范围
    plt.xlabel('Truth_x')
    plt.ylabel('Predict_x')
    plt.title('The comparison of truth_x and predict_x')
    # plt.legend(loc='upper right')
    plt.show()
